
from isaaclab.utils import configclass
from isaaclab.sensors.frame_transformer.frame_transformer_cfg import FrameTransformerCfg
from isaaclab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg 
from isaaclab.assets import  AssetBaseCfg, RigidObjectCfg
from isaaclab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg 
from isaaclab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg, UsdFileCfg
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR

from arm_lab.assets.config.micro_ur import MICRO_UR_CFG
from arm_lab.tasks.manager_based.reach.reach_env_cfg import ReachEnvCfg
from arm_lab.tasks.manager_based.reach import mdp
import math

@configclass
class MicroUrEnvCfg(ReachEnvCfg):

    tool_link_name = "tool_link"
    joints_name = ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6"]
    
    def __post_init__(self):
        # post init of parent
        super().__post_init__()

        # ------------------------------Sence------------------------------
        # switch robot 
        self.scene.robot = MICRO_UR_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")

        # ------------------------------Commands----------------------------------- 
        self.commands.ee_pose.body_name = self.tool_link_name

        # ------------------------------Actions-----------------------------------
        self.actions.arm_action = mdp.JointPositionActionCfg(
            asset_name="robot", joint_names=self.joints_name, scale=0.5, use_default_offset=True, preserve_order=True
        )

        # ------------------------------Observations------------------------------
        self.observations.policy.joint_pos.params["asset_cfg"].joint_names = self.joints_name
        self.observations.policy.joint_vel.params["asset_cfg"].joint_names = self.joints_name

        # ------------------------------Events------------------------------------
        self.events.randomize_rigid_body_material = None
        self.events.randomize_apply_external_force_torque = None
        self.events.randomize_reset_joints.params["position_range"] = (0.75, 1.25)
        # self.events.randomize_rigid_body_mass = None
        # self.events.randomize_rigid_body_inertia = None

        # ------------------------------Rewards------------------------------
        self.rewards.end_effector_position_tracking.params["asset_cfg"].body_names = [self.tool_link_name]
        self.rewards.end_effector_position_tracking_fine_grained.params["asset_cfg"].body_names = [self.tool_link_name]
        self.rewards.end_effector_orientation_tracking.params["asset_cfg"].body_names = [self.tool_link_name]

        # tracking reward
        self.rewards.end_effector_position_tracking.weight = -0.2
        self.rewards.end_effector_position_tracking_fine_grained.weight = 0.1
        self.rewards.end_effector_orientation_tracking.weight = -0.15

        # action penalty
        self.rewards.action_rate.weight = -1e-4
        self.rewards.joint_vel.weight = -1e-4

        # torque penalty
        self.rewards.applied_torque_limits.weight = -1.0e-4
        self.rewards.joint_torques_l2.weight = -2.5e-05

        # joints limits
        self.rewards.joint_pos_limits.weight = -1.0
        self.rewards.joint_vel_limits.weight = 0.0

        # contact penalty
        self.rewards.undesired_contacts.weight = -5.0


        # ------------------------------Terminations------------------------------

        # ------------------------------Curriculum------------------------------
        # self.curriculum.action_rate = None
        # self.curriculum.joint_vel = None
        
        # set rewards which weight 0.0 to None
        # If the weight of rewards is 0, set rewards to None
        if self.__class__.__name__ == "MicroUrEnvCfg":
            self.disable_zero_weight_rewards()
        
